#include "HuffmanTree.h"
#include <string>
#include <climits>

HuffmanTree::HuffmanTree(const std::map<unsigned char, long long>& weights) {
    for (auto it = weights.begin(); it != weights.end(); ++it) {
        unsigned char character = it->first;
        long long weight = it->second;
        tree.emplace_back(character, weight);
    }
}


// 从前 n 个节点中选择两个权重最小且未被选中的节点
void HuffmanTree::selectTwo(int n, int& s1, int& s2) {
    long long minWeight = INT_MAX;
    s1 = s2 = -1;
    for (int i = 0; i < n; ++i) {
        if (tree[i].father == 0 && tree[i].weight < minWeight) {
            minWeight = tree[i].weight;
            s1 = i;
        }
    }
    minWeight = INT_MAX;
    for (int i = 0; i < n; ++i) {
        if (tree[i].father == 0 && tree[i].weight < minWeight && i != s1) {
            minWeight = tree[i].weight;
            s2 = i;
        }
    }
}

// 构建哈夫曼树并生成编码
void HuffmanTree::generateTree() {
    int n = tree.size();
    for (int i = n; i < 2 * n - 1; ++i) {
        int s1 = 0, s2 = 0;
        selectTwo(i, s1, s2);
        tree[s1].father = i;
        tree[s2].father = i;
        tree.emplace_back(0, tree[s1].weight + tree[s2].weight); // 新建父节点
        tree.back().lc = s1;
        tree.back().rc = s2;
    }

    // 生成哈夫曼编码
    for (int i = 0; i < n; ++i) {
        std::string code = "";
        for (int current = i, parent = tree[i].father; parent;
            current = parent, parent = tree[current].father) {
            code = (tree[parent].lc == current ? '0' : '1') + code;
        }
        encodingMap[tree[i].ch] = code;
    }
}

// 返回编码映射
const std::map<unsigned char, std::string>& HuffmanTree::getEncodingMap() const {
    return encodingMap;
}

// 返回 Huffman 树
const std::vector<HuffmanNode>& HuffmanTree::getTree() const {
    return tree;
}
